package limiter

import (
	"sync"
	"time"

	"golang.org/x/time/rate"
)

type IPRateLimiter struct {
	ips    map[string]*rate.Limiter
	mu     sync.RWMutex
	rate   rate.Limit
	burst  int
	ttl    time.Duration
	ticker *time.Ticker
}

func NewIPRateLimiter(r rate.Limit, b int, ttl time.Duration) *IPRateLimiter {
	limiter := &IPRateLimiter{
		ips:   make(map[string]*rate.Limiter),
		rate:  r,
		burst: b,
		ttl:   ttl,
	}

	// Start cleanup routine
	limiter.ticker = time.NewTicker(ttl)
	go limiter.cleanup()

	return limiter
}

func (l *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
	l.mu.Lock()
	defer l.mu.Unlock()

	limiter, exists := l.ips[ip]
	if !exists {
		limiter = rate.NewLimiter(l.rate, l.burst)
		l.ips[ip] = limiter
	}

	return limiter
}

func (l *IPRateLimiter) cleanup() {
	for range l.ticker.C {
		l.mu.Lock()
		for ip, limiter := range l.ips {
			if limiter.Allow() {
				delete(l.ips, ip)
			}
		}
		l.mu.Unlock()
	}
}

func (l *IPRateLimiter) Stop() {
	if l.ticker != nil {
		l.ticker.Stop()
	}
}

type RateLimiter interface {
	Allow() bool
	AllowN(time.Time, int) bool
}
